import pygame
import sys
import random

# 初始化pygame
pygame.init()

# 设置屏幕尺寸
screen_width = 640
screen_height = 480

# 创建窗口
screen = pygame.display.set_mode((screen_width, screen_height))

# 设置标题
pygame.display.set_caption("Snake Game")

# 设置颜色
white = (255, 255, 255)
green = (0, 255, 0)
red = (255, 0, 0)
black = (0, 0, 0)

# 蛇的初始位置和速度
snake_pos = [100, 50]
snake_body = [[100, 50], [90, 50], [80, 50]]

# 果实的生成位置
food_pos = [random.randrange(1, (screen_width // 10)) * 10,
            random.randrange(1, (screen_height // 10)) * 10]
food_spawn = True

# 移动方向
direction = 'RIGHT'
change_to = direction

# 游戏结束标志
game_over = False

# 游戏速度
clock = pygame.time.Clock()

# 游戏主循环
while not game_over:

    # 处理事件
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            game_over = True
        elif event.type == pygame.KEYDOWN:
            if event.key == pygame.K_UP:
                change_to = 'UP'
            if event.key == pygame.K_DOWN:
                change_to = 'DOWN'
            if event.key == pygame.K_LEFT:
                change_to = 'LEFT'
            if event.key == pygame.K_RIGHT:
                change_to = 'RIGHT'

    # 更改方向
    if change_to == 'UP' and direction != 'DOWN':
        direction = 'UP'
    if change_to == 'DOWN' and direction != 'UP':
        direction = 'DOWN'
    if change_to == 'LEFT' and direction != 'RIGHT':
        direction = 'LEFT'
    if change_to == 'RIGHT' and direction != 'LEFT':
        direction = 'RIGHT'

    # 根据方向移动蛇头
    if direction == 'UP':
        snake_pos[1] -= 10
    if direction == 'DOWN':
        snake_pos[1] += 10
    if direction == 'LEFT':
        snake_pos[0] -= 10
    if direction == 'RIGHT':
        snake_pos[0] += 10

    # 增加蛇的身体长度
    snake_body.insert(0, list(snake_pos))

    # 如果蛇吃掉果实
    if snake_pos[0] == food_pos[0] and snake_pos[1] == food_pos[1]:
        food_spawn = False
    else:
        snake_body.pop()

    # 如果食物不存在则生成新的食物
    if not food_spawn:
        food_pos = [random.randrange(1, (screen_width // 10)) * 10,
                    random.randrange(1, (screen_height // 10)) * 10]
    food_spawn = True

    # 清屏
    screen.fill(black)

    # 绘制蛇的身体
    for pos in snake_body:
        pygame.draw.rect(screen, green, pygame.Rect(pos[0], pos[1], 10, 10))

    # 绘制果实
    pygame.draw.rect(screen, white, pygame.Rect(food_pos[0], food_pos[1], 10, 10))

    # 刷新屏幕
    pygame.display.flip()

    # 检测是否碰壁
    if snake_pos[0] < 0 or snake_pos[0] > screen_width - 10:
        game_over = True
    if snake_pos[1] < 0 or snake_pos[1] > screen_height - 10:
        game_over = True

    # 检测蛇体碰撞
    for block in snake_body[1:]:
        if snake_pos[0] == block[0] and snake_pos[1] == block[1]:
            game_over = True

    # 控制游戏速度
    clock.tick(15)

# 结束pygame
pygame.quit()
sys.exit()